# -*- coding:utf-8 -*-

import datetime

from sqlalchemy import Column, BigInteger, DateTime, String, text, JSON
from sqlalchemy import func, and_, or_
from sqlalchemy.future import select
from sqlalchemy import create_engine
from sqlalchemy.orm import sessionmaker
from sqlalchemy.ext.asyncio import create_async_engine
from sqlalchemy.ext.asyncio import async_sessionmaker, AsyncSession
from sqlalchemy.ext.asyncio import AsyncAttrs
from sqlalchemy.orm import DeclarativeBase
from sqlalchemy.ext.declarative import as_declarative, declared_attr

from urllib.parse import quote_plus
from config.env import DataBaseConfig
from utils.snowflake import generator

from utils.logger import logger
from utils.dependence import PageFilter

ASYNC_SQLALCHEMY_DATABASE_URL = (
    f'mysql+asyncmy://{DataBaseConfig.db_username}:{quote_plus(DataBaseConfig.db_password)}@'
    f'{DataBaseConfig.db_host}:{DataBaseConfig.db_port}/{DataBaseConfig.db_database}'
)
if DataBaseConfig.db_type == 'postgresql':
    ASYNC_SQLALCHEMY_DATABASE_URL = (
        f'postgresql+asyncpg://{DataBaseConfig.db_username}:{quote_plus(DataBaseConfig.db_password)}@'
        f'{DataBaseConfig.db_host}:{DataBaseConfig.db_port}/{DataBaseConfig.db_database}'
    )

if DataBaseConfig.db_type == 'sqlite':
    ASYNC_SQLALCHEMY_DATABASE_URL = (
        f'sqlite:///{DataBaseConfig.db_database}.db'
    )
    
    engine = create_engine(
        ASYNC_SQLALCHEMY_DATABASE_URL, connect_args={"check_same_thread": False}, echo=DataBaseConfig.db_echo
    )
    SessionLocal = sessionmaker(autocommit=False, autoflush=False, bind=engine)
else:
    async_engine = create_async_engine(
        ASYNC_SQLALCHEMY_DATABASE_URL,
        echo=DataBaseConfig.db_echo,
        max_overflow=DataBaseConfig.db_max_overflow,
        pool_size=DataBaseConfig.db_pool_size,
        pool_recycle=DataBaseConfig.db_pool_recycle,
        pool_timeout=DataBaseConfig.db_pool_timeout,
    )
    AsyncSessionLocal = async_sessionmaker(autocommit=False, autoflush=False, bind=async_engine)

@as_declarative()
class Base(AsyncAttrs):
    id = Column(BigInteger, primary_key=True, default=generator.generate_id)
    extra = Column(JSON, default=dict, comment='扩展字段')
    remark =  Column(String(256), comment='备注')
    create_source =  Column(String(256), comment='创建来源')
    create_time = Column(DateTime, nullable=False, default=datetime.datetime.now, comment='创建时间')
    create_user =  Column(BigInteger, default=-1, comment='创建人')
    update_time = Column(DateTime, nullable=False, default=datetime.datetime.now, onupdate=datetime.datetime.now, index=True, comment='更新时间')
    update_user =  Column(BigInteger, default=-1, comment='更新人')
    del_flag = Column(String(1), nullable=False, default='0', server_default=text("'0'"), comment='删除标志(0代表存在 1代表删除)')

    def __repr__(self):
        return "[%s%s]" % (self.__tablename__, str({c.name: getattr(self, c.name) for c in self.__table__.columns}))

    def single_to_dict(self):
        return {c.name: getattr(self, c.name).strftime("%Y-%m-%d %H:%M:%S") if isinstance(getattr(self, c.name), datetime) else getattr(self, c.name) for c in self.__table__.columns}


class BaseMixin:
    """model的基类,所有model都必须继承"""
    id = Column(BigInteger, primary_key=True, default=generator.generate_id)
    extra = Column(JSON, default=dict, comment='扩展字段')
    remark =  Column(String(256), comment='备注')
    create_source =  Column(String(256), comment='创建来源')
    create_time = Column(DateTime, nullable=False, default=datetime.datetime.now, comment='创建时间')
    create_user =  Column(BigInteger, default=-1, comment='创建人')
    update_time = Column(DateTime, nullable=False, default=datetime.datetime.now, onupdate=datetime.datetime.now, index=True, comment='更新时间')
    update_user =  Column(BigInteger, default=-1, comment='更新人')
    del_flag = Column(String(1), nullable=False, default='0', server_default=text("'0'"), comment='删除标志(0代表存在 1代表删除)')

    def __repr__(self):
        return "[%s%s]" % (self.__tablename__, str({c.name: getattr(self, c.name) for c in self.__table__.columns}))

    def single_to_dict(self):
        return {c.name: getattr(self, c.name).strftime("%Y-%m-%d %H:%M:%S") if isinstance(getattr(self, c.name), datetime) else getattr(self, c.name) for c in self.__table__.columns}

    @classmethod
    def filter(cls, page_filter: PageFilter):
        filter = [getattr(cls, k) == v for k, v in page_filter.filter.items()] if page_filter.filter else []
        in_ = [getattr(cls, k).in_(v) for k, v in page_filter.in_.items()] if page_filter.in_ else []
        like = [getattr(cls, k).like(f'%{v}%') for k, v in page_filter.like.items()] if page_filter.like else []
        order_by = [getattr(cls, k).asc() if v else getattr(cls, k).desc() for k, v in page_filter.order_by.items()] if page_filter.order_by else []
        return (filter, in_, like, order_by)

    @classmethod
    async def page(cls, page_no, page_size, page_filter: PageFilter, session: AsyncSession):
        filter, in_, like, order_by = cls.filter(page_filter)
        total_count = await cls.count(session, page_filter)
        stm = select(cls).filter(cls.del_flag == '0', *filter, *in_, *like)
        if page_size == -1:
            page_size = total_count
            page_no = 1
        result = await session.execute(stm.order_by(*order_by).offset((page_no - 1) * page_size).limit(page_size))
        # 计算总页数
        total_pages = (total_count // page_size) + (1 if total_count % page_size > 0 else 0)

        return result.scalars().all(), total_count, total_pages

    @classmethod
    async def get_filter(cls, session: AsyncSession, page_filter: PageFilter = None, and_or: bool = True):
        filter, in_, like, _= cls.filter(page_filter)
        result = await session.execute(selectselect(cls).filter(
            and_(cls.del_flag == '0', 
                and_(*filter, *in_, *like) if and_or else or_(*filter, *in_, *like)
            )))
        return result.scalar(),all()

    @classmethod
    async def count(cls, session: AsyncSession, page_filter: PageFilter = None, and_or: bool = True):
        filter, in_, like, _= cls.filter(page_filter)
        total_count = await session.execute(select(func.count(getattr(cls, 'id'))).filter(
            and_(cls.del_flag == '0', 
                and_(*filter, *in_, *like) if and_or else or_(*filter, *in_, *like)
            )))
        return total_count.scalar()

    @classmethod
    async def first(cls, obj, session: AsyncSession):
        filter = [getattr(cls, f'{column.name}') == getattr(obj, f'{column.name}') for column in cls.__table__.columns if hasattr(obj, f'{column.name}') and getattr(obj, f'{column.name}') is not None]
        result = await session.execute(select(cls).filter(cls.del_flag == '0', *filter))
        return result.scalars().first()
    
    @classmethod
    async def get_ref(cls, session: AsyncSession, ref_ids: list[int], ref_table_name: str, ref_type: str):
        ref_table = get_model_class_by_table_name(ref_table_name)
        stm = select(TbAssociation.source_id, ref_table).join(ref_table, and_(ref_table.id == TbAssociation.target_id)).where(TbAssociation.del_flag == '0', TbAssociation.association_type == ref_type, TbAssociation.source_id.in_(ref_ids))
        result = await session.execute(stm)
        result = result.all()

        return {id: [item[1] for item in result if item[0] == id] for id in ref_ids}

    async def add(self, session: AsyncSession, commit: bool=False, flush: bool = True):
        session.add(self)
        if commit:
            await session.commit()
        elif flush:
            await session.flush()
        return self

    async def add_ref(self, session: AsyncSession, ref_ids: list[int], ref_table_name: str, ref_type: str, create_user: int = -1, commit: bool=False, flush: bool = True):
        ref_table = get_model_class_by_table_name(ref_table_name)
        for ref_id in ref_ids:
            await TbAssociation(source_id=self.id, target_id=ref_id, association_type=ref_type, create_user=create_user).add(session, commit=commit, flush=flush)
        if commit:
            await session.commit()
        elif flush:
            await session.flush()

    async def delete_ref(self, session: AsyncSession, ref_ids: list[int], ref_table_name: str, ref_type: str, delete_user: int = -1, commit: bool=False, flush: bool = True):
        ref_table = get_model_class_by_table_name(ref_table_name)
        for ref_id in ref_ids:
            db_data = await TbAssociation.first(TbAssociation(source_id=self.id, target_id=ref_id, association_type=ref_type, del_flag=0), session)
            db_data.update_user = delete_user
            await db_data.delete(session, commit=commit, flush=flush)
        if commit:
            await session.commit()
        elif flush:
            await session.flush()

    async def delete(self, session: AsyncSession, commit: bool=False, flush: bool = True):
        self.del_flag = 1
        if commit:
            await session.commit()
        elif flush:
            await session.flush()
        return self
    
    async def update(self, session: AsyncSession, data: dict={}, commit: bool=False, flush: bool = True):
        for key, val in data.items():
            if val is None or key not in self.__table__.columns: continue
            setattr(self, key, val)
        if commit:
            await session.commit()
        elif flush:
            await session.flush()
        return self



async def get_db():
    """
    每一个请求处理完毕后会关闭当前连接，不同的请求使用不同的连接
    :return:
    """
    if DataBaseConfig.db_type == 'sqlite':
        current_db = SessionLocal()
        try:
            yield current_db
            current_db.commit()
        finally:
            current_db.close()
    else:
        async with AsyncSessionLocal() as current_db:
            try:
                yield current_db
                await current_db.commit()
            except:
                await current_db.rollback()
                raise
            finally:
                await current_db.close()


async def init_create_table():
    """
    应用启动时初始化数据库连接
    :return:
    """
    logger.info('初始化数据库连接...')
    if DataBaseConfig.db_type == 'sqlite':
        Base.metadata.create_all(bind=engine)
    else:
        async with async_engine.begin() as conn:
            await conn.run_sync(Base.metadata.create_all)
    logger.info('数据库连接成功')

# 根据表名获取模型类的函数
def get_model_class_by_table_name(table_name: str):
    for model_class in Base.registry._class_registry.values():
        if hasattr(model_class, '__tablename__') and model_class.__tablename__ == table_name:
            return model_class
    return None  # 如果没有找到匹配的模型类，返回 None


class TbAssociation(Base):
    __tablename__ = 'tb_association'

    id = Column(BigInteger, primary_key=True)
    source_id = Column(BigInteger)
    target_id = Column(BigInteger)
    association_type = Column(String(128))